Knowledge bases for Amazon Bedrock でメタデータフィルタリングを試してみる

秘匿レベル：全社

# オフィスの利用方法マニュアル

## 1. 入退室について

- オフィスの利用可能時間は、平日の 9:00 から 18:00 までです。
- 入室の際は、受付で名前と所属をお伝えください。
- 退室時は、必ず施錠し、鍵を受付に返却してください。

{
    "nextToken": "string",
    "retrievalResults": [
        {
            "content": {
                "text": "string"
            },
            "location": {
                "s3Location": {
                    "uri": "string"
                },
                "type": "S3"
            },
            "metadata": {
                "string": {...}|[...]|123|123.4|"string"|True|None
            },
            "score": 123.0
        },
    ]
}

takakuni% tree .
.
├── 就業規則.md
├── 就業規則.md.metadata.json
├── オフィスの利用方法.md
└── オフィスの利用方法.md.metadata.json

{
	"metadataAttributes": {
		"target": "company-wide",
		"year": 2024,
		"for_managers": true
	}
}

ALTER TABLE bedrock_integration.bedrock_kb
ADD COLUMN IF NOT EXISTS target VARCHAR(100),
ADD COLUMN IF NOT EXISTS year SMALLINT,
ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;

takakuni % tree .
.
├── company-wide
│   ├── 就業規則.md
│   ├── 会議室マニュアル.md
│   └── オフィスの利用方法.md
└── department
    ├── engineer
    │   ├── 個人検証用AWSアカウントのセットアップ方法.md
    │   └── 目標設定について.md
    └── sales
        ├── 営業資料.pdf
        └── 目標設定について.md

import os
import boto3

bucket_name = os.environ.get("BUCKET_NAME")
region = os.environ.get("AWS_REGION")
kb_id = os.environ.get("KNOWLEDGE_BASE_ID")
attribute = os.environ.get("USER_ATTRIBUTE")
client = boto3.client("bedrock-agent-runtime", region_name=region)

response = client.retrieve(
    knowledgeBaseId=kb_id,
    retrievalConfiguration={
        "vectorSearchConfiguration": {
            "filter": {
                "startsWith": {
                    "key": "x-amz-bedrock-kb-source-uri",
                    "value": f"s3://{bucket_name}/department/{attribute}/",
                }
            },
            "numberOfResults": 1,
        },
    },
    retrievalQuery={"text": "個人目標の設定方法を教えてください"},
)

company-wide
├── 2023年度事業計画.md
├── 2023年度事業計画.md.metadata.json
├── 2024年度事業計画.md
├── 2024年度事業計画.md.metadata.json
├── 就業規則.md
├── 就業規則.md.metadata.json
├── 会議室マニュアル.md
├── 会議室マニュアル.md.metadata.json
├── オフィスの利用方法.md
└── オフィスの利用方法.md.metadata.json

{
	"metadataAttributes": {
		"target": "company-wide",
		"year": 2024,
		"for_managers": true
	}
}

export REGION="us-west-2"
export CLUSTER_ARN="Aurora クラスター ARN"
export SECRET_ARN="シークレット ARN"
export DATABASE_NAME="postgresql"

aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "ALTER TABLE bedrock_integration.bedrock_kb \
  ADD COLUMN IF NOT EXISTS target VARCHAR(100), \
  ADD COLUMN IF NOT EXISTS year SMALLINT, \
  ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;"

[cloudshell-user@ip-10-130-63-143 ~]$ export REGION="us-west-2"
[cloudshell-user@ip-10-130-63-143 ~]$ export CLUSTER_ARN=arn:aws:rds:us-west-2:XXXXXXXXXXXX:cluster:kb-metadata-vctrdb-cluster
[cloudshell-user@ip-10-130-63-143 ~]$ export SECRET_ARN=arn:aws:secretsmanager:us-west-2:XXXXXXXXXXXX:secret:kb-metadata-kb-vctrdb-secret-XXXXXX
[cloudshell-user@ip-10-130-63-143 ~]$ export DATABASE_NAME="postgresql"
[cloudshell-user@ip-10-130-63-143 ~]$ aws rds-data --region $REGION execute-statement \
>   --resource-arn $CLUSTER_ARN \
>   --secret-arn $SECRET_ARN \
>   --database $DATABASE_NAME \
>   --sql "ALTER TABLE bedrock_integration.bedrock_kb \
>   ADD COLUMN IF NOT EXISTS target VARCHAR(100), \
>   ADD COLUMN IF NOT EXISTS year SMALLINT, \
>   ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;"
{
    "numberOfRecordsUpdated": 0,
    "generatedFields": []
}

import os
import boto3

bucket_name = os.environ.get("BUCKET_NAME")
region = os.environ.get("AWS_REGION")
kb_id = os.environ.get("KNOWLEDGE_BASE_ID")
client = boto3.client("bedrock-agent-runtime", region_name=region)

filter = {
    "andAll": [
        {"equals": {"key": "year", "value": 2024}},
        {"equals": {"key": "for_managers", "value": True}},
    ]
}

response = client.retrieve_and_generate(
    input={"text": "来年はどんな事業に注力する予定ですか？"},
    retrieveAndGenerateConfiguration={
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": kb_id,
            "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "filter": filter,
                    "numberOfResults": 5,
                }
            },
        },
        "type": "KNOWLEDGE_BASE",
    },
)

print(response["output"]["text"])

(app-py3.12) takakuni@app % python medatada_filter.py
株式会社サンプルは2024年度に、宇宙旅行事業とAI搭載ロボットペット事業の2つの新規事業に注力する予定です。 宇宙旅行事業では、月面でのリゾート「ルナ・パラダイス」の建設、宇宙飛行士養成プログラムの提供、宇宙食の開発と販売を行います。AI搭載ロボットペット事業では、感情認識AIを搭載した「エモ・ペット」の開発、セレブリティとのコラボレーションモデルの販売、ロボットペット用アクセサリーの開発と販売を行う予定です。